package com.aabte.lota.auth.auth.shiro;

import com.aabte.lota.auth.auth.token.JWTTokenUtils;
import com.aabte.lota.auth.auth.token.exception.TokenException;
import com.aabte.lota.auth.auth.token.payload.PermissionDTO;
import com.aabte.lota.auth.auth.token.payload.RoleDTO;
import com.aabte.lota.auth.auth.token.payload.TokenPayloadDTO;
import com.google.common.collect.Sets;
import java.util.Set;
import java.util.stream.Collectors;
import javax.annotation.Resource;
import org.apache.shiro.authc.AuthenticationException;
import org.apache.shiro.authc.AuthenticationInfo;
import org.apache.shiro.authc.AuthenticationToken;
import org.apache.shiro.authc.SimpleAuthenticationInfo;
import org.apache.shiro.authz.AuthorizationInfo;
import org.apache.shiro.authz.SimpleAuthorizationInfo;
import org.apache.shiro.realm.AuthorizingRealm;
import org.apache.shiro.subject.PrincipalCollection;

/**
 * @author Daniel
 * @version 1.0
 * @date 2020/4/5
 */
public class JwtTokenRealm extends AuthorizingRealm {

  private static final String REALM_NAME = "jwtTokenRealm";
  @Resource private JWTTokenUtils jwtTokenUtils;

  @Override
  public boolean supports(AuthenticationToken token) {
    return token instanceof JwtToken;
  }

  @Override
  protected AuthorizationInfo doGetAuthorizationInfo(PrincipalCollection principals) {

    JwtToken jwtToken = (JwtToken) principals.getPrimaryPrincipal();

    TokenPayloadDTO tokenPayload;
    try {
      tokenPayload = jwtTokenUtils.parseToken(jwtToken.getToken());
    } catch (TokenException e) {
      throw new AuthenticationException(e);
    }

    Set<String> stringPermissions;
    if (null != tokenPayload.getPermissions()) {
      stringPermissions =
          tokenPayload.getPermissions().stream()
              .map(PermissionDTO::getName)
              .collect(Collectors.toSet());
    } else {
      stringPermissions = Sets.newHashSet();
    }

    Set<String> stringRoles;
    if (null != tokenPayload.getRoles()) {
      stringRoles =
          tokenPayload.getRoles().stream().map(RoleDTO::getName).collect(Collectors.toSet());
    } else {
      stringRoles = Sets.newHashSet();
    }

    SimpleAuthorizationInfo simpleAuthorizationInfo = new SimpleAuthorizationInfo();
    simpleAuthorizationInfo.setStringPermissions(stringPermissions);
    simpleAuthorizationInfo.setRoles(stringRoles);
    return simpleAuthorizationInfo;
  }

  @Override
  protected AuthenticationInfo doGetAuthenticationInfo(AuthenticationToken token)
      throws AuthenticationException {
    JwtToken jwtToken = (JwtToken) token;

    TokenPayloadDTO tokenPayload;
    try {
      tokenPayload = jwtTokenUtils.parseToken(jwtToken.getToken());
    } catch (TokenException e) {
      throw new AuthenticationException(e.getMessage(), e);
    }

    if (null != tokenPayload) {
      return new SimpleAuthenticationInfo(token, token, REALM_NAME);
    }
    throw new AuthenticationException("tokenPayload is null");
  }
}
